package org.goobs.csp;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Stack;

public class CSP {
	private HashSet<ICSPArc<?, ?>> arcs = new HashSet<ICSPArc<?, ?>>();

	@SuppressWarnings("unchecked")
	private HashMap<Integer, CSPVariable> variables = new HashMap<Integer, CSPVariable>();

	// complex relation support
	private boolean registered = false;
	private HashSet<IConstraint> constraints;

	// --------------------
	// Constructors
	// --------------------
	public CSP() {
		this(false);
	}

	private CSP(boolean registered) {
		this.registered = registered;
		if (!registered) {
			constraints = new HashSet<IConstraint>();
		}
	}

	// --------------------
	// Setup Functions
	// --------------------

	public void addVariable(CSPVariable<?> newVar) {
		if (registered) {
			throw new IllegalStateException(
					"CSP cannot be modified when finding solutions");
		}
		if (newVar.totalArcCount() != 0) {
			throw new IllegalArgumentException(
					"CSP Variable has links when being added to CSP: " + newVar);
		}
		if (variables.containsKey(newVar.uniqueNum())) {
			throw new IllegalArgumentException(
					"CSP already contains variable: " + newVar);
		}
		variables.put(newVar.uniqueNum(), newVar);
	}

	@SuppressWarnings("unchecked")
	private void addArc(ICSPArc newArc) {
		if (registered) {
			throw new IllegalStateException(
					"CSP cannot be modified when finding solutions");
		}
		CSPVariable<?> tail = newArc.getTail();
		CSPVariable<?> head = newArc.getHead();
		if (!variables.containsKey(head.uniqueNum())
				|| !variables.containsKey(tail.uniqueNum())) {
			throw new IllegalArgumentException(
					"adding an arc to CSP without end variables present in CSP");
		} else if (arcs.contains(newArc)) {
			throw new IllegalArgumentException(
					"adding an arc to CSP which already exists");
		} else {
			arcs.add(newArc);
			tail.addArc(newArc);
			head.addArc(newArc);
		}
	}

	@SuppressWarnings("unchecked")
	public boolean removeVariable(CSPVariable toRemove) {
		if (registered) {
			throw new IllegalStateException(
					"CSP cannot be modified when finding solutions");
		}
		return variables.remove(toRemove.uniqueNum()) != null;
	}

	@SuppressWarnings("unchecked")
	private boolean removeArc(ICSPArc toRemove) {
		if (registered) {
			throw new IllegalStateException(
					"CSP cannot be modified when finding solutions");
		}
		CSPVariable tail = toRemove.getTail();
		CSPVariable head = toRemove.getHead();
		if (!variables.containsKey(head.uniqueNum())
				|| !variables.containsKey(tail.uniqueNum())) {
			return false;
		}
		if (!tail.removeArc(toRemove) || !head.removeArc(toRemove)) {
			// the csp is inconsistent
			throw new IllegalStateException(
					"inconsistent CSP state. Did you remove the same arc twice?");
		}
		return arcs.remove(toRemove);
	}

	// --------------------
	// Complex Relation Support
	// --------------------

	@SuppressWarnings("unchecked")
	protected Iterator<CSPVariable> variableIterator() {
		return (Iterator<CSPVariable>) variables.values().iterator();
	}

	public void addConstraint(IConstraint cons) {
		constraints.add(cons);
	}

	public void clearConstraints() {
		constraints = new HashSet<IConstraint>();
	}

	// --------------------
	// Algorithm Functions
	// --------------------

	@SuppressWarnings("unchecked")
	public synchronized boolean register() {
		if (this.registered) {
			return false;
		}

		// --Set constraints
		for (IConstraint cons : constraints) {
			Collection<ICSPArc> arcs = cons.generateArcs(this);
			for (ICSPArc arc : arcs) {
				this.addArc(arc);
			}
		}

		// --Reset variables
		for (CSPVariable var : variables.values()) {
			var.reset();
		}

		this.registered = true;
		return true;
	}

	public synchronized void release() {
		this.registered = false;
		arcs = new HashSet<ICSPArc<?, ?>>();

	}

	public CSP copyCSP() {
		CSP rtn = new CSP(this.registered);
		rtn.arcs = this.arcs;
		rtn.variables = new HashMap<Integer, CSPVariable>();
		Iterator<Integer> iter = this.variables.keySet().iterator();
		while (iter.hasNext()) {
			Integer next = iter.next();
			rtn.variables.put(next, this.variables.get(next).cloneValue());
		}
		return rtn;
	}

	// @SuppressWarnings("unchecked")
	// private void ac3(){
	// //Build the queue
	// Queue <ICSPArc> arcQueue = new LinkedList <ICSPArc> ();
	// for(ICSPArc arc : arcs){
	// arcQueue.add(arc);
	// }
	//		
	// //--Remove inconsistent values
	// while(!arcQueue.isEmpty()){
	// ICSPArc arc = arcQueue.poll();
	// if(removeInconsistentValues(arc)){
	// Iterator <ICSPArc> toAdd = arc.getTail().inboundArcs();
	// while(toAdd.hasNext()){
	// arcQueue.add(toAdd.next());
	// }
	// }
	// }
	//		
	// //--Free some memory
	// varMap = null;
	// }
	//	
//	@SuppressWarnings("unchecked")
//	private boolean removeInconsistentValues(ICSPArc arc){
//		//--State variables
//		ICSPVariable tail = localCopy(arc.getTail());
//		ICSPVariable baseTail = arc.getTail();
//		Object tailVal = baseTail.getValue();
//		ICSPVariable head = localCopy(arc.getHead());
//		ICSPVariable baseHead = arc.getHead();
//		Object headVal = baseHead.getValue();
//		//--Algorithm variables
//		boolean removed = false;
//		Iterator tailVals = tail.acceptableValues();
//		//--For every value of tail
//		while(tailVals.hasNext()){
//			//If there is no value of head, remove that value from tail
//			baseTail.setValue(tailVals.next());
//			Iterator headVals = head.acceptableValues();
//			boolean exists = false;
//			//--Look through values of HEAD
//			//(special case: head == tail)
//			if(head == tail){
//				exists = arc.isOK();
//			}else{
//				while(headVals.hasNext()){
//					baseHead.setValue(headVals.next());
//					if(arc.isOK()){
//						exists = true;
//					}
//				}
//			}
//			//--If something existed, remove it
//			if(!exists){
//				tailVals.remove();
//				removed = true;
//				println("rejected value " + tail.getValue() + " from var " + tail + "\t\t
//						arc: " + arc);
//							println(((CSPVariable) tail).acceptableValues);
//			}
//		}
//		baseTail.setValue(tailVal);
//		baseHead.setValue(headVal);
//		if(removed){
//			println(detailedString());
//		}
//		return removed;
//	}

	private class Diff {

		private HashMap <CSPVariable, ArrayList<Object>> changes = new HashMap<CSPVariable, ArrayList<Object>>();
		private HashMap <CSPVariable, Object> allocated = new HashMap <CSPVariable, Object> ();
		private CSPVariable curVar;
		private Iterator<Object> valIter;

		private Diff(CSPVariable curVal, Iterator<Object> valIter) {
			this.curVar = curVal;
			this.valIter = valIter;
		}

		private void registerDomainLost(CSPVariable var, Object domain) {
			ArrayList<Object> lst = changes.get(var);
			if (lst == null) {
				lst = new ArrayList<Object>();
				changes.put(var, lst);
			}
			lst.add(domain);
		}
		
		private void registerVariableAssigned(CSPVariable var){
			registerVariableAssigned(var, var.getValue());
		}
		
		private void registerVariableAssigned(CSPVariable var, Object oldVal){
			if(oldVal != null){
				throw new IllegalStateException();
			}
			allocated.put(var, oldVal);
		}

		private void restore(CSP csp) {
			// --Restore variable domains
			for( CSPVariable var : changes.keySet() ){
				ArrayList<Object> vals = changes.get(var);
				for (Object val : vals) {
					csp.variables.get(var.uniqueNum()).addAcceptableValue(val);
				}
			}
			// --Restore variable values
			for( CSPVariable var : allocated.keySet() ){
//				println("--------reseting var " + var + " to " + allocated.get(var));
				csp.variables.get(var.uniqueNum()).unchooseValue(allocated.get(var));
			}
			// --Restore iterator states
			csp.valIter = valIter;
			csp.chosenVar = curVar;
		}
	}

	private String prefix = "   ";
	public boolean findSolution() {
		boolean stop = false;
		boolean skipOverhead = false;
		boolean rtn = false;
		println("Start");
		while (!stop) {
			int state = nextSolution(skipOverhead, prefix);
			switch (state) {
			case FAIL:
				prefix = prefix.substring(3);
				if (dfsStack.isEmpty()) {
					rtn = false;
					stop = true;
				} else {
					println(prefix + "<--Up");
					skipOverhead = true;
					upSync();
				}
				break;
			case BASE_CASE:
				println(prefix + "BASE");
				upSync(); // (base case is not a real state)
				rtn = true;
				stop = true;
				break;
			case RECURSE:
				println(prefix + "-->Recurse");
				prefix += "   ";
				skipOverhead = false;
				downSync();
				break;
			default:
				throw new IllegalStateException(
						"Unknown return type from nextSolution(): " + state);
			}
		}
		println("End");
		return rtn;
	}

	private void downSync() {
		Diff lastState = new Diff(this.chosenVar, this.valIter);
		dfsStack.push(lastState);
		this.chosenVar = null;
		this.valIter = null;
	}

	private void upSync() {
		Diff lastState = dfsStack.pop();
		lastState.restore(this);
	}

	private Iterator<Object> valIter;
	private CSPVariable chosenVar;
	private Stack<Diff> dfsStack = new Stack<Diff>();

	private static final int FAIL = -1;
	private static final int BASE_CASE = 1;
	private static final int RECURSE = 2;

	private int nextSolution(boolean skipOverhead, String prefix) {
		if (!skipOverhead) {
			if (!registered) {
				throw new IllegalStateException(
						"CSP is not registered. Cannot find solutions");
			}

			// --Run arc consistency
			// (run ac3)
			ac3();
			// (ensure that every variable has at least one acceptable value)
			Iterator<CSPVariable> iter = this.variableIterator();
			while (iter.hasNext()) {
				CSPVariable<?> var = iter.next();
				if (var.numAcceptableValues() == 0 || !var.isConsistent()) {
					println(prefix + "fail ac3 on var: " + var);
					return FAIL;
				}
			}

			// --Check consistency
			boolean consistent = true;
			for (ICSPArc arc : arcs) {
				if (!arc.isOK()) {
					consistent = false;
					break;
				}
			}
			if (!consistent) {
				println(prefix + "fail consistency");
				return FAIL;
			}

			// --Chose a variable to check
			// (variables)
			if (chosenVar == null) {
				int mostConstrained = Integer.MAX_VALUE;
				iter = this.variableIterator();
				// (find unassigned variable)
				while (iter.hasNext()) {
					CSPVariable<?> var = iter.next();
					if (var.getValue() == null) {
						// (set the most constrained variable first)
						if (var.numAcceptableValues() < mostConstrained) {
							chosenVar = var;
						}
					}
				}
				// (check for base case)
				if (chosenVar == null) {
					return BASE_CASE; // BASE CASE: No unassigned variables
				}
				// (create the value iterator)
				valIter = chosenVar.acceptableValues();
			}
		}
		
		// --Skip Overhead Starts Here
		while (valIter.hasNext()) {
			// (recursive case)
			//(register the variable assignment)
			if(!dfsStack.isEmpty()){
				dfsStack.peek().registerVariableAssigned(chosenVar, null);
			}
			//(assign the variable)
			Object val = valIter.next();
			chosenVar.chooseValue(val);
			//(return the command to recurse)
			println(prefix + chosenVar + " : " + val);
			return RECURSE;
		}
		
		println(prefix + "fail no more values");
		return FAIL;
	}

	private void ac3() {
		Queue <ICSPArc> arcs = new LinkedList <ICSPArc> ();
		for( ICSPArc arc : this.arcs ){
			arcs.offer(arc);
		}
		
		while(!arcs.isEmpty()){
			ICSPArc arc = arcs.poll();
			if(removeInconsistentValues(arc)){
				Iterator <ICSPArc> iter = arc.getTail().inboundArcs();
				while(iter.hasNext()){
					arcs.offer(iter.next());
				}
			}
		}
	}
	
	private boolean removeInconsistentValues(ICSPArc arc){
		// --Variables and save values
		boolean removed = false;
		CSPVariable tail = arc.getTail();
		CSPVariable head = arc.getHead();
		Object tailSv = tail.getValue();
		Object headSv = head.getValue();
		
		// --For each val in domain(tail)
		Iterator <Object> domain = tail.acceptableValues();
		while(domain.hasNext()){
			//(get vals)
			Object tailObj = domain.next();
			tail.setValue(tailObj);
			Iterator <Object> otherDom = head.acceptableValues();
			// --If no variable in domain(head) satisfies...
			boolean good = false;
			if(head == tail){
				good = arc.isOK();
			}else{
				while(otherDom.hasNext()){
					Object headObj = otherDom.next();
					head.setValue(headObj);
					if(arc.isOK()){
						good = true;
					}
				}
			}
			// --Then remove that val from domain(tail)
			if(!good){
				//(register loss)
				if(!dfsStack.isEmpty()){
					dfsStack.peek().registerDomainLost(tail, tailObj);
				}
				//(remove value)
				domain.remove();
				//(removed is now true)
				removed = true;
			}
		}
		
		tail.setValue(tailSv);
		head.setValue(headSv);
		return removed;
	}

	private void println(Object o) {
//		System.out.println(o);
	}

	public void printState() {
		System.out.println(toString());
	}

	public String toString() {
		String rtn = "";
		Iterator<CSPVariable> iter = variables.values().iterator();
		while (iter.hasNext()) {
			CSPVariable<?> term = iter.next();
			rtn = rtn + term + "\n";
		}
		return rtn;
	}

	@SuppressWarnings("unchecked")
	public String detailedString() {
		String rtn = "";
		Iterator<CSPVariable> iter = variables.values().iterator();
		while (iter.hasNext()) {
			CSPVariable term = iter.next();
			String vals = "";
			Iterator vIt = term.acceptableValues();
			while (vIt.hasNext()) {
				vals = vals + vIt.next().toString() + ", ";
			}
			rtn = rtn + term + "  accept vals: {" + vals + "}\n";
		}
		return rtn;
	}

	@SuppressWarnings("unchecked")
	public void printArcs() {
		for (ICSPArc arc : arcs) {
			System.out.println(arc);
		}
	}

}
